#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge
import depthai as dai
import numpy as np

class AlignedDepthPublisher(Node):
    def __init__(self):
        super().__init__('aligned_depth_publisher')
        
        # ROS2 Publishers
        self.rgb_pub = self.create_publisher(Image, 'camera/aligned_rgb', 10)
        self.depth_pub = self.create_publisher(Image, 'camera/aligned_depth', 10)
        self.camera_info_pub = self.create_publisher(CameraInfo, 'camera/info', 10)
        self.bridge = CvBridge()

        # DepthAI初始化
        try:
            self.pipeline = self.create_pipeline()
            self.device = dai.Device()
            self.device.startPipeline(self.pipeline)
        except Exception as e:
            self.get_logger().error(f"Device initialization failed: {e}")
            raise

        # 创建输出队列
        self.rgb_queue = self.device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
        self.depth_queue = self.device.getOutputQueue(name="depth", maxSize=4, blocking=False)

        # 相机参数
        self.calib_data = self.device.readCalibration()
        self.camera_info = self.get_camera_info()

    def create_pipeline(self):
        pipeline = dai.Pipeline()

        # RGB摄像头配置
        cam_rgb = pipeline.create(dai.node.ColorCamera)
        cam_rgb.setBoardSocket(dai.CameraBoardSocket.CAM_A)
        cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
        cam_rgb.setFps(30)
        cam_rgb.setIspScale(2, 3)  # 输出分辨率 1280x720
        #cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)

        # 立体深度配置
        mono_left = pipeline.create(dai.node.MonoCamera)
        mono_right = pipeline.create(dai.node.MonoCamera)
        stereo = pipeline.create(dai.node.StereoDepth)

        # Mono相机参数
        mono_res = dai.MonoCameraProperties.SensorResolution.THE_480_P
        mono_left.setResolution(mono_res)
        mono_left.setBoardSocket(dai.CameraBoardSocket.CAM_B)
        mono_right.setResolution(mono_res)
        mono_right.setBoardSocket(dai.CameraBoardSocket.CAM_C)

        # 深度对齐配置
        stereo.setDepthAlign(dai.CameraBoardSocket.CAM_A)
        #stereo.setOutputSize(mono_left.getResolutionWidth(), mono_left.getResolutionHeight())
        
        # 深度处理参数
        stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
        stereo.setLeftRightCheck(True)
        stereo.setSubpixel(False)
        #stereo.setExtendedDisparity(False)

        # 深度优化配置
        stereo.initialConfig.setMedianFilter(dai.MedianFilter.KERNEL_7x7)  # 中值滤波
        
        # 后处理配置
        config = stereo.initialConfig.get()
        # 散斑滤波器
        config.postProcessing.speckleFilter.enable = True
        config.postProcessing.speckleFilter.speckleRange = 50
        # 时间滤波器
        config.postProcessing.temporalFilter.enable = True
        config.postProcessing.temporalFilter.persistencyMode = dai.StereoDepthConfig.PostProcessing.TemporalFilter.PersistencyMode.VALID_8_OUT_OF_8
        # 空间滤波器
        config.postProcessing.spatialFilter.enable = True
        config.postProcessing.spatialFilter.holeFillingRadius = 2
        config.postProcessing.spatialFilter.numIterations = 1
        # 阈值滤波器
        config.postProcessing.thresholdFilter.minRange = 400   # 40cm
        config.postProcessing.thresholdFilter.maxRange = 15000 # 15m
        config.postProcessing.decimationFilter.decimationFactor = 1
        stereo.initialConfig.set(config)

        # 节点连接
        mono_left.out.link(stereo.left)
        mono_right.out.link(stereo.right)

        # 输出配置
        xout_rgb = pipeline.createXLinkOut()
        xout_rgb.setStreamName("rgb")
        cam_rgb.isp.link(xout_rgb.input)

        xout_depth = pipeline.createXLinkOut()
        xout_depth.setStreamName("depth")
        stereo.depth.link(xout_depth.input)

        return pipeline

    def get_camera_info(self):
        """生成符合ROS规范的相机标定信息"""
        info = CameraInfo()
        info.header.frame_id = "camera_frame"
        info.width = 1280
        info.height = 720
        
        # 获取内参矩阵
        cam_matrix = self.calib_data.getCameraIntrinsics(
            dai.CameraBoardSocket.CAM_A, 
            dai.Size2f(info.width, info.height)
        )
        
        # 构造K矩阵
        info.k = [
            cam_matrix[0][0], 0.0, cam_matrix[0][2],
            0.0, cam_matrix[1][1], cam_matrix[1][2],
            0.0, 0.0, 1.0
        ]
        
        # 畸变系数
        distortion_coeffs = self.calib_data.getDistortionCoefficients(dai.CameraBoardSocket.CAM_A)
        info.d = [float(x) for x in distortion_coeffs]
        
        info.p = [
            cam_matrix[0][0], 0.0, cam_matrix[0][2], 0.0,
            0.0, cam_matrix[1][1], cam_matrix[1][2], 0.0,
            0.0, 0.0, 1.0, 0.0
        ]
        info.distortion_model = "plumb_bob"
        return info

    def publish_data(self):
        # 获取最新数据
        rgb_data = self.rgb_queue.get()
        depth_data = self.depth_queue.get()

        if rgb_data is not None and depth_data is not None:
            current_time = self.get_clock().now().to_msg()
            
            # 处理RGB图像
            frame_rgb = rgb_data.getCvFrame()
            rgb_msg = self.bridge.cv2_to_imgmsg(frame_rgb, "bgr8")
            rgb_msg.header.stamp = current_time
            rgb_msg.header.frame_id = "camera_frame"
            self.rgb_pub.publish(rgb_msg)

            # 处理深度图像（单位：毫米）
            frame_depth = depth_data.getFrame().astype(np.uint16)
            depth_msg = self.bridge.cv2_to_imgmsg(frame_depth, "16UC1")
            depth_msg.header.stamp = current_time
            depth_msg.header.frame_id = "camera_frame"
            self.depth_pub.publish(depth_msg)

            # 发布相机参数
            self.camera_info.header.stamp = current_time
            self.camera_info_pub.publish(self.camera_info)

def main(args=None):
    rclpy.init(args=args)
    try:
        node = AlignedDepthPublisher()
        while rclpy.ok():
            node.publish_data()
            rclpy.spin_once(node, timeout_sec=0.001)
    except Exception as e:
        print(f"Error: {e}")
    finally:
        if 'node' in locals():
            node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
